#!/usr/bin/env python
import sys
import warnings
import yaml
import os
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Optional

from datetime import datetime

from investment_team.crew import InvestmentTeam
from investment_team.llm_config import prediction_llm

warnings.filterwarnings("ignore", category=SyntaxWarning, module="pysbd")

# 这个主文件旨在让您在本地运行您的
# 团队，因此请避免在此文件中添加不必要的逻辑。
# 替换为您想要测试的输入，它将自动
# 插入任何任务和代理信息

def load_fund_list():
    """
    从配置文件加载基金列表。
    
    Returns:
        list: 基金列表，每个元素包含 'fund' 字段
        
    Raises:
        FileNotFoundError: 配置文件不存在
        yaml.YAMLError: YAML 格式错误
        KeyError: 配置文件格式不正确
    """
    # 获取配置文件路径（相对于当前文件）
    config_path = Path(__file__).parent / 'config' / 'funds.yaml'
    
    try:
        with open(config_path, 'r', encoding='utf-8') as f:
            config = yaml.safe_load(f)
        
        # 验证配置文件格式
        if not config or 'funds' not in config:
            raise KeyError("配置文件格式错误：缺少 'funds' 键")
        
        if not isinstance(config['funds'], list):
            raise KeyError("配置文件格式错误：'funds' 必须是列表")
        
        # 验证每个基金项是否包含必要字段
        for idx, fund in enumerate(config['funds']):
            if not isinstance(fund, dict):
                raise KeyError(f"配置文件格式错误：第 {idx + 1} 个基金项必须是字典")
            if 'fund' not in fund:
                raise KeyError(f"配置文件格式错误：第 {idx + 1} 个基金项缺少 'fund' 字段")
        
        return config['funds']
        
    except FileNotFoundError:
        raise FileNotFoundError(f"配置文件不存在：{config_path}")
    except yaml.YAMLError as e:
        raise yaml.YAMLError(f"YAML 格式错误：{e}")

def _process_single_fund(fund: dict, model_name: str) -> tuple[str, Optional[Exception]]:
    """
    处理单个基金，返回基金标识和可能的异常。
    
    Args:
        fund: 基金信息字典，包含 'fund' 字段
        model_name: 模型名称
        
    Returns:
        tuple: (基金标识, 异常对象)，如果成功则异常为 None
    """
    # 提取基金代码用于标识（从 fund['fund'] 中提取，格式如 'sz159570,港股通创新药ETF'）
    fund_identifier = fund['fund'].split(',')[0].strip() if ',' in fund['fund'] else fund['fund']
    
    try:
        inputs = { 
            'current_time': str(datetime.now()),
            'fund' : fund['fund'],
            'model_name': model_name,
        }
        InvestmentTeam().crew().kickoff(inputs=inputs)
        return fund_identifier, None
    except Exception as e:
        return fund_identifier, e

def run():
    """
    并发运行团队。
    使用线程池并发处理所有基金，根据 CPU 核心数自动设置并发数。
    收集所有错误并在最后统一报告。
    """
    # 从配置文件加载基金列表
    try:
        fundList = load_fund_list()
    except (FileNotFoundError, yaml.YAMLError, KeyError) as e:
        raise Exception(f"加载基金配置失败：{e}")

    if not fundList:
        print("基金列表为空，无需处理")
        return

    # 从 LLM 对象获取模型名称
    model_name = prediction_llm.model if hasattr(prediction_llm, 'model') else ''
    
    # 对于 I/O 密集型任务（大模型 API 调用），线程数应该是 CPU 核心数的倍数
    # 因为线程在等待 I/O 时会被阻塞，不占用 CPU，可以设置更多线程
    cpu_count = os.cpu_count() or 1
    # 设置为 CPU 核心数的 3 倍，适合 I/O 密集型任务
    base_workers = cpu_count * 3
    # 设置合理范围：最小 5，最大 20（避免过多并发导致资源耗尽或触发 API 限流）
    max_workers = min(max(base_workers, 5), 20)
    # 如果基金数量较少，可以适当增加并发，但不超过基金数量
    if len(fundList) < max_workers:
        max_workers = len(fundList)
    
    print(f"并发线程数设置为: {max_workers} (CPU 核心数: {cpu_count}, 基金数量: {len(fundList)})")
    
    # 收集所有错误
    errors = []
    successful_funds = []
    
    # 使用线程池并发执行
    with ThreadPoolExecutor(max_workers=max_workers) as executor:
        # 提交所有任务
        future_to_fund = {
            executor.submit(_process_single_fund, fund, model_name): fund 
            for fund in fundList
        }
        
        # 收集结果
        for future in as_completed(future_to_fund):
            fund = future_to_fund[future]
            try:
                fund_identifier, exception = future.result()
                if exception is None:
                    successful_funds.append(fund_identifier)
                    print(f"✓ 基金 {fund_identifier} 处理成功")
                else:
                    errors.append((fund_identifier, exception))
                    print(f"✗ 基金 {fund_identifier} 处理失败: {exception}")
            except Exception as e:
                # 如果 future.result() 本身抛出异常
                fund_identifier = fund['fund'].split(',')[0].strip() if ',' in fund['fund'] else fund['fund']
                errors.append((fund_identifier, e))
                print(f"✗ 基金 {fund_identifier} 处理失败: {e}")
    
    # 统一报告结果
    total_count = len(fundList)
    success_count = len(successful_funds)
    error_count = len(errors)
    
    print(f"\n=== 处理完成 ===")
    print(f"总计: {total_count} 个基金")
    print(f"成功: {success_count} 个")
    print(f"失败: {error_count} 个")
    
    # 如果有错误，抛出包含所有错误的异常
    if errors:
        error_messages = []
        for fund_identifier, exception in errors:
            error_messages.append(f"基金 {fund_identifier}: {str(exception)}")
        
        error_summary = "\n".join(error_messages)
        raise Exception(f"部分基金处理失败 ({error_count}/{total_count}):\n{error_summary}")


def train():
    """
    在给定迭代次数下训练团队。
    """
    inputs = {
        "topic": "AI LLMs",
        'current_year': str(datetime.now().year)
    }
    try:
        InvestmentTeam().crew().train(n_iterations=int(sys.argv[1]), filename=sys.argv[2], inputs=inputs)

    except Exception as e:
        raise Exception(f"An error occurred while training the crew: {e}")

def replay():
    """
    从特定任务重放团队执行。
    """
    try:
        InvestmentTeam().crew().replay(task_id=sys.argv[1])

    except Exception as e:
        raise Exception(f"An error occurred while replaying the crew: {e}")

def test():
    """
    测试团队执行并返回结果。
    """
    inputs = {
        "topic": "AI LLMs",
        "current_year": str(datetime.now().year)
    }
    
    try:
        InvestmentTeam().crew().test(n_iterations=int(sys.argv[1]), eval_llm=sys.argv[2], inputs=inputs)

    except Exception as e:
        raise Exception(f"An error occurred while testing the crew: {e}")
 
if __name__ == "__main__":
    run()